import os
import json
import logging
import argparse
from datetime import datetime
from sqlalchemy import create_engine, text
from sqlalchemy.orm import sessionmaker
from sqlalchemy.ext.declarative import declarative_base

# 配置日志
logging.basicConfig(
  level=logging.INFO,
  format='%(asctime)s - %(levelname)s - %(message)s',
  handlers=[
    logging.FileHandler('db_initialization.log', encoding='utf-8'),
    logging.StreamHandler()
  ]
)
logger = logging.getLogger(__name__)

# 数据库配置
DATABASE_URI = 'sqlite:///your_database.db'
SQL_DIRECTORY = "sql_scripts"  # SQL脚本目录
RETRY_STATE_FILE = "init_retry_state.json"  # 记录重试状态的文件
MAX_RETRY_COUNT = 3  # 最大重试次数

# 安全配置
DANGEROUS_COMMANDS = {'DROP', 'TRUNCATE', 'DELETE', 'ALTER'}

# ORM基础类
Base = declarative_base()


class DatabaseInitializer:
  def __init__(self, app_filter):
    # 初始化数据库连接（ORM方式）
    self.engine = create_engine(DATABASE_URI)
    self.Session = sessionmaker(bind=self.engine)
    self.session = self.Session()

    # 初始化状态
    self.app_filter = app_filter
    self.success_tables = set()
    self.empty_tables = set()
    self.failed_files = self.load_retry_state()  # 加载上次失败的文件
    self.retry_count = {f: 0 for f in self.failed_files}  # 记录每个文件的重试次数

    try:
      logger.info("=" * 50)
      logger.info(f"开始数据库初始化（{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}）")
      logger.info(f"目标应用: {app_filter} | 最大重试次数: {MAX_RETRY_COUNT}")
      logger.warning("⚠️ 安全模式：仅执行INSERT/CREATE，不删除任何数据")

      # 核心流程
      self.check_database_connection()
      self.check_tables_data()
      self.process_sql_files()  # 处理文件（优先重试失败的）
      self.save_retry_state()  # 保存本次失败的文件

    except Exception as e:
      logger.error(f"初始化流程中断: {str(e)}", exc_info=True)
    finally:
      self.session.close()
      logger.info("数据库连接已关闭")
      logger.info(f"初始化流程结束（剩余未完成文件: {len(self.failed_files)}）")
      logger.info("=" * 50 + "\n")

  def check_database_connection(self):
    """检查数据库连接（ORM方式）"""
    try:
      Base.metadata.create_all(self.engine, checkfirst=True)  # 确保表结构存在
      logger.info(f"成功连接到数据库: {DATABASE_URI}")
    except Exception as e:
      logger.error(f"数据库连接失败: {str(e)}")
      raise

  def check_tables_data(self):
    """检查所有表的数据状态"""
    try:
      # 通过ORM反射获取所有表
      Base.metadata.reflect(bind=self.engine)
      all_tables = Base.metadata.tables.keys()
      logger.info(f"通过ORM检测到 {len(all_tables)} 个表...")

      for table in all_tables:
        if table.startswith('sqlite_'):
          continue
        # 应用过滤
        if self.app_filter != 'all' and not table.startswith(f'{self.app_filter}_'):
          continue
        # 检查数据量
        count = self.session.query(text(f"SELECT COUNT(*) FROM {table}")).scalar()
        if count > 0:
          self.success_tables.add(table)
          logger.info(f"表 [{table}] 已有数据（{count}条），跳过处理")
        else:
          self.empty_tables.add(table)
          logger.info(f"表 [{table}] 为空，需要初始化")

    except Exception as e:
      logger.error(f"表数据检查失败: {str(e)}", exc_info=True)
      raise

  def process_sql_files(self):
    """处理SQL文件（优先重试失败的，再处理新文件）"""
    if not os.path.exists(SQL_DIRECTORY):
      logger.error(f"SQL目录不存在: {SQL_DIRECTORY}")
      raise FileNotFoundError(f"SQL directory {SQL_DIRECTORY} not found")

    # 获取所有SQL文件并排序
    all_sql_files = [f for f in os.listdir(SQL_DIRECTORY) if f.endswith(".sql")]
    all_sql_files.sort()

    # 分离：需要重试的文件 和 新文件
    retry_files = [f for f in self.failed_files if f in all_sql_files]
    new_files = [f for f in all_sql_files if f not in self.failed_files]

    logger.info(f"\n待处理文件统计: 重试文件 {len(retry_files)} 个 | 新文件 {len(new_files)} 个")

    # 1. 先处理重试文件
    if retry_files:
      logger.info("\n===== 开始处理重试文件 =====")
      for file in retry_files:
        self._process_single_file(file, is_retry=True)

    # 2. 再处理新文件
    if new_files:
      logger.info("\n===== 开始处理新文件 =====")
      for file in new_files:
        self._process_single_file(file, is_retry=False)

  def _process_single_file(self, sql_file, is_retry=False):
    """处理单个SQL文件"""
    # 检查是否超过最大重试次数
    if is_retry and self.retry_count.get(sql_file, 0) >= MAX_RETRY_COUNT:
      logger.error(f"文件 {sql_file} 已达最大重试次数（{MAX_RETRY_COUNT}），放弃处理")
      return

    # 提取应用和表名
    file_app = 'rbac' if sql_file.startswith('rbac_') else 'web' if sql_file.startswith('web_') else None
    table_name = self._extract_table_name(sql_file)

    # 过滤条件检查
    if self.app_filter != 'all' and file_app != self.app_filter:
      logger.info(f"跳过非目标应用文件: {sql_file}（应用: {file_app}）")
      return

    # 已有数据的表直接跳过
    if table_name and table_name in self.success_tables:
      logger.info(f"表 [{table_name}] 已有数据，跳过文件: {sql_file}")
      # 从失败列表中移除（如果存在）
      if sql_file in self.failed_files:
        self.failed_files.remove(sql_file)
      return

    # 执行文件
    file_path = os.path.join(SQL_DIRECTORY, sql_file)
    logger.info(f"\n处理文件: {sql_file}（重试次数: {self.retry_count.get(sql_file, 0)}/{MAX_RETRY_COUNT}）")

    try:
      with open(file_path, 'r', encoding='utf-8') as f:
        sql_content = f.read()

      # 分割并执行SQL语句
      statements = [s.strip() for s in sql_content.split(';') if s.strip() and not s.startswith('--')]
      success = True

      for i, stmt in enumerate(statements, 1):
        # 安全检查
        if any(cmd in stmt.upper() for cmd in DANGEROUS_COMMANDS):
          logger.error(f"拒绝执行危险语句 {i}: {stmt[:30]}...")
          success = False
          continue

        if not (stmt.upper().startswith('INSERT') or stmt.upper().startswith('CREATE')):
          logger.warning(f"跳过非初始化语句 {i}: {stmt[:30]}...")
          continue

        # 执行语句（ORM会话）
        try:
          self.session.execute(text(stmt))
          self.session.commit()
          logger.info(f"成功执行语句 {i}/{len(statements)}")
        except Exception as e:
          self.session.rollback()
          logger.error(f"语句 {i} 执行失败: {str(e)}")
          success = False
          break  # 单条语句失败则中断当前文件

      if success:
        logger.info(f"✅ 文件 {sql_file} 执行成功")
        if sql_file in self.failed_files:
          self.failed_files.remove(sql_file)
        if table_name:
          self.success_tables.add(table_name)
      else:
        raise Exception("部分语句执行失败")

    except Exception as e:
      logger.error(f"❌ 文件 {sql_file} 执行失败: {str(e)}")
      if sql_file not in self.failed_files:
        self.failed_files.append(sql_file)
      self.retry_count[sql_file] = self.retry_count.get(sql_file, 0) + 1

  def _extract_table_name(self, sql_file):
    """从文件名提取表名（假设格式: {app}_{table}_xxx.sql）"""
    parts = sql_file.split('_')
    if len(parts) >= 2 and parts[0] in ['rbac', 'web']:
      table_parts = []
      for part in parts[1:]:
        if part.isdigit() or '-' in part:
          break
        table_parts.append(part)
      return f"{parts[0]}_{'_'.join(table_parts)}" if table_parts else None
    return None

  def load_retry_state(self):
    """加载上次失败的文件列表"""
    if os.path.exists(RETRY_STATE_FILE):
      try:
        with open(RETRY_STATE_FILE, 'r', encoding='utf-8') as f:
          return json.load(f)
      except Exception as e:
        logger.warning(f"加载重试状态失败: {str(e)}，将重新开始")
    return []

  def save_retry_state(self):
    """保存本次失败的文件列表"""
    try:
      with open(RETRY_STATE_FILE, 'w', encoding='utf-8') as f:
        json.dump(self.failed_files, f, ensure_ascii=False, indent=2)
      logger.info(f"已保存重试状态（{len(self.failed_files)} 个文件需要重试）")
    except Exception as e:
      logger.error(f"保存重试状态失败: {str(e)}")


if __name__ == "__main__":
  parser = argparse.ArgumentParser(description='ORM数据库初始化（支持自动重试）')
  parser.add_argument('--app', type=str, default='all',
                      help='指定应用（rbac/web/all）')
  args = parser.parse_args()

  # 启动初始化
  DatabaseInitializer(app_filter=args.app)
